package top.bulk.ratelimit.handle;

import com.google.common.util.concurrent.RateLimiter;
import org.springframework.stereotype.Component;
import top.bulk.ratelimit.constant.RateLimiterConst;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

/**
 * 基于Guava的 RateLimiter 实现的限流，注意，这是一种 平滑限流 的方案
 *
 * <b>这种限流不支持分布式！只支持本地限流</b>
 *
 * @author 散装java
 * @date 2024-06-26
 */
@Component(RateLimiterConst.GUAVA_RATELIMITER)
// 抑制 @Beta 警告
@SuppressWarnings("UnstableApiUsage")
public class GuavaRateLimiterHandler implements RateLimiterHandler {
    /**
     * 用于存储 RateLimiter 实例的映射。
     */
    static final Map<String, RateLimiter> RATE_LIMITER_MAP = new ConcurrentHashMap<>();

    @Override
    public boolean limit(String key, Long max, Long time, TimeUnit unit) {

        // 获取 RateLimiter 实例，如果不存在则创建，如果配置变更，则返回新的
        RateLimiter limiter = RATE_LIMITER_MAP.compute(key, (k, existingLimiter) -> {
            // 计算速率，即是 QPS
            double rate = max.doubleValue() / unit.toSeconds(time);
            // 如果不存在或速率不匹配，则创建新的RateLimiter实例；
            if (existingLimiter == null) {
                return RateLimiter.create(rate);
            } else if (existingLimiter.getRate() != rate) {
                // 如果速率不同，设置新的速率
                existingLimiter.setRate(rate);
                return existingLimiter;
            } else {
                // 如果速率匹配，则直接返回现有的 RateLimiter
                return existingLimiter;
            }
        });

        return !limiter.tryAcquire();
    }
}
